#include <gtest/gtest.h>

#include <memory>

#include "./cuda/cuda_vector.hpp"
#include "./hash_join.hpp"
#include "scheduler.hpp"
#include "task_set.hpp"

TEST(SchedulerTest, Init) {
  // spdlog::set_level(spdlog::level::debug);
  Scheduler scheduler;
  scheduler.start();

  std::vector<int> data(10000, 1);  // NOLINT

  // operator
  auto process_func = [](std::vector<int>::iterator begin,
                         std::vector<int>::iterator end,
                         int multiplier) -> int {
    int sum = 0;
    for (auto it = begin; it != end; ++it) {
      sum += (*it) * multiplier;
    }
    return sum;
  };

  // finalizer
  auto finalize_func =
      [](std::vector<int>&& partial_results) -> int {  // NOLINT
    int total = 0;
    for (auto& result : partial_results) {
      total += result;
    }
    return total;
  };

  // make task set
  for (int iter = 0; iter <= 1000; ++iter) {  // NOLINT
    auto resource_group = std::make_unique<ResourceGroup>();

    auto task_set =
        make_task_set(data.begin(), data.end(), process_func, finalize_func, 2);

    // result future
    auto future = task_set->get_future();

    // get resource group
    resource_group->add_task_set(std::move(task_set));

    // submit resource group
    size_t slot = scheduler.submit_resource_group(std::move(resource_group));

    int result = future.get();
    EXPECT_EQ(result, 20000);
  }
}

TEST(SchedulerTest, MultiResourceGroup) {
  // spdlog::set_level(spdlog::level::debug);
  Scheduler scheduler;
  scheduler.start();

  std::vector<int> data(1000, 1);  // NOLINT

  // operator
  auto process_func = [](std::vector<int>::iterator begin,
                         std::vector<int>::iterator end,
                         int multiplier) -> int {
    int sum = 0;
    for (auto it = begin; it != end; ++it) {
      sum += (*it) * multiplier;
    }
    return sum;
  };

  // finalizer
  auto finalize_func =
      [](std::vector<int>&& partial_results) -> int {  // NOLINT
    int total = 0;
    for (auto& result : partial_results) {
      total += result;
    }
    return total;
  };

  for (int iter = 0; iter < 100; ++iter) {  // NOLINT
    std::vector<std::future<int>> futures;

    for (int i = 0; i < MAX_RESOURCE_GROUPS; ++i) {
      auto task_set = make_task_set(data.begin(), data.end(), process_func,
                                    finalize_func, i);
      futures.push_back(task_set->get_future());

      auto resource_group = std::make_unique<ResourceGroup>();
      resource_group->add_task_set(std::move(task_set));
      scheduler.submit_resource_group(std::move(resource_group));
    }

    for (int i = 0; i < MAX_RESOURCE_GROUPS; ++i) {
      EXPECT_EQ(futures[i].get(), 1000 * i);
    }
  }
}

TEST(SchedulerTest, ChunkSize) {
  spdlog::set_level(spdlog::level::info);
  Scheduler scheduler;
  scheduler.start();

  const int DATANUM = 10000'0000;

  std::vector<int> data(DATANUM, 1);  // NOLINT

  // operator
  auto process_func = [](std::vector<int>::iterator begin,
                         std::vector<int>::iterator end,
                         int multiplier) -> int {
    int sum = 0;
    for (auto it = begin; it != end; ++it) {
      sum += (*it) * multiplier;
    }
    return sum;
  };

  // finalizer
  auto finalize_func =
      [](std::vector<int>&& partial_results) -> int {  // NOLINT
    int total = 0;
    for (auto& result : partial_results) {
      total += result;
    }
    return total;
  };

  // make task set
  auto resource_group = std::make_unique<ResourceGroup>();

  auto task_set =
      make_task_set(data.begin(), data.end(), process_func, finalize_func, 2);

  // result future
  auto future = task_set->get_future();

  // get resource group
  resource_group->add_task_set(std::move(task_set));

  // submit resource group
  auto start_time = std::chrono::steady_clock::now();
  size_t slot = scheduler.submit_resource_group(std::move(resource_group));
  int result = future.get();
  auto end_time = std::chrono::steady_clock::now();

  EXPECT_EQ(result, DATANUM * 2);
  std::chrono::duration<double, std::milli> execution_time =
      end_time - start_time;
  double milliseconds = execution_time.count();
  std::cout << milliseconds << "\n";

  // run in sequence
  start_time = std::chrono::steady_clock::now();
  process_func(data.begin(), data.end(), 2);
  end_time = std::chrono::steady_clock::now();

  execution_time = end_time - start_time;
  milliseconds = execution_time.count();
  std::cout << milliseconds << "\n";

  scheduler.export_analyse("data1.bin");
}

TEST(SchedulerTest, FutureParams) {
  spdlog::set_level(spdlog::level::info);
  Scheduler scheduler;
  scheduler.start();

  const int DATANUM = 10000'0000;

  std::vector<int> data(DATANUM, 1);  // NOLINT

  // operator
  auto process_func = [](std::vector<int>::iterator begin,
                         std::vector<int>::iterator end,
                         int multiplier) -> int {
    int sum = 0;
    for (auto it = begin; it != end; ++it) {
      sum += (*it) * multiplier;
    }
    return sum;
  };

  // finalizer
  auto finalize_func =
      [](std::vector<int>&& partial_results) -> int {  // NOLINT
    int total = 0;
    for (auto& result : partial_results) {
      total += result;
    }
    return total;
  };

  auto process_func2 = [](std::vector<int>::iterator begin,
                          std::vector<int>::iterator end, int result) {
    for (auto it = begin; it != end; ++it) {
      *it = result;
    }
  };

  auto resource_group = std::make_unique<ResourceGroup>();

  auto task_set1 =
      make_task_set(data.begin(), data.end(), process_func, finalize_func, 2);
  auto future1 = task_set1->get_future();

  auto task_set2 = make_task_set(data.begin(), data.end(), process_func2,
                                 VOID_FINALIZE, std::move(future1));
  auto future2 = task_set2->get_future();

  // get resource group
  resource_group->add_task_set(std::move(task_set1));
  resource_group->add_task_set(std::move(task_set2));

  // submit resource group
  auto start_time = std::chrono::steady_clock::now();
  size_t slot = scheduler.submit_resource_group(std::move(resource_group));
  bool result = future2.get();
  auto end_time = std::chrono::steady_clock::now();

  EXPECT_EQ(result, true);
  std::chrono::duration<double, std::milli> execution_time =
      end_time - start_time;
  double milliseconds = execution_time.count();
  std::cout << milliseconds << "\n";

  for (int i = 0; i < 10; ++i) {  // NOLINT
    EXPECT_EQ(data[i], DATANUM * 2);
  }

  scheduler.export_analyse("data2.bin");
}

TEST(SchedulerTest, HashJoin) {
  using namespace hashjoin;
  // build query plan
  Plan plan;

  Table table1 = {{1, "a", 4},   // NOLINT
                  {2, "b", 5},   // NOLINT
                  {3, "c", 6}};  // NOLINT
  Table table2 = {{1, 1.1F},     // NOLINT
                  {2, 2.2F},     // NOLINT
                  {4, 3.3F}};    // NOLINT

  // print_table(table1);
  // print_table(table2);

  plan.new_input(std::move(table1));
  plan.new_input(std::move(table2));

  // projection of table1
  std::vector<std::tuple<size_t, DataType>> scan1_attrs = {
      {0, DataType::INT32}, {1, DataType::STRING}};  // get col 0 and col 1
  size_t scan1 = plan.new_scan_node(0, scan1_attrs);

  // projection of table2
  std::vector<std::tuple<size_t, DataType>> scan2_attrs = {
      {0, DataType::INT32}, {1, DataType::FLOAT}};  // get col 0 and col 1
  size_t scan2 = plan.new_scan_node(1, scan2_attrs);

  // hash join of table1 and table2
  std::vector<std::tuple<size_t, DataType>> join_attrs = {
      {0, DataType::INT32}, {1, DataType::STRING}, {3, DataType::FLOAT}};
  size_t join = plan.new_join_node(scan1, scan2, 0, 0, join_attrs);

  plan.root = join;

  // execute with scheduler
  Scheduler scheduler;
  scheduler.start();

  auto rg = std::make_unique<ResourceGroup>();

  auto ret = generate_resource_group(rg.get(), plan, plan.root);
  auto finish = rg->get_future();

  scheduler.submit_resource_group(std::move(rg));

  finish.get();

  Table ground_truth = {{1, "a", 1.1F},   // NOLINT
                        {2, "b", 2.2F}};  // NOLINT
  // print_table(result);

  EXPECT_EQ(ground_truth, *ret);
}

TEST(SchedulerTest, HashJoinLarge) {
  using namespace hashjoin;
  Plan plan;

  // NOLINTBEGIN
  Table table1;
  for (int i = 0; i < 10'000; i++) {
    std::string str_val = "str_" + std::to_string(i);
    int val3 = 100 + i;
    table1.push_back({i, str_val, val3});
  }

  Table table2;
  for (int i = 0; i < 8'000; i++) {
    float float_val = i * 1.5f;
    table2.push_back({i, float_val});
  }

  plan.new_input(std::move(table1));
  plan.new_input(std::move(table2));

  std::vector<std::tuple<size_t, DataType>> scan1_attrs = {
      {0, DataType::INT32}, {1, DataType::STRING}};
  size_t scan1 = plan.new_scan_node(0, scan1_attrs);

  std::vector<std::tuple<size_t, DataType>> scan2_attrs = {
      {0, DataType::INT32}, {1, DataType::FLOAT}};
  size_t scan2 = plan.new_scan_node(1, scan2_attrs);

  std::vector<std::tuple<size_t, DataType>> join_attrs = {
      {0, DataType::INT32}, {1, DataType::STRING}, {3, DataType::FLOAT}};
  size_t join = plan.new_join_node(scan1, scan2, 0, 0, join_attrs);

  plan.root = join;

  Scheduler scheduler;
  scheduler.start();

  auto rg = std::make_unique<ResourceGroup>();

  auto ret = generate_resource_group(rg.get(), plan, plan.root);
  auto finish = rg->get_future();

  const int DATANUM = 10000'0000;

  std::vector<int> data(DATANUM, 1);  // NOLINT
  // operator
  auto process_func = [](std::vector<int>::iterator begin,
                         std::vector<int>::iterator end,
                         int multiplier) -> int {
    int sum = 0;
    for (auto it = begin; it != end; ++it) {
      sum += (*it) * multiplier;
    }
    return sum;
  };

  // finalizer
  auto finalize_func =
      [](std::vector<int>&& partial_results) -> int {  // NOLINT
    int total = 0;
    for (auto& result : partial_results) {
      total += result;
    }
    return total;
  };

  std::future<int> fut;

  {
    // make task set
    auto resource_group = std::make_unique<ResourceGroup>();

    auto task_set =
        make_task_set(data.begin(), data.end(), process_func, finalize_func, 2);

    // result future
    fut = task_set->get_future();

    // get resource group
    resource_group->add_task_set(std::move(task_set));

    scheduler.submit_resource_group(std::move(resource_group));
  }

  auto start_time = std::chrono::steady_clock::now();
  scheduler.submit_resource_group(std::move(rg));
  finish.get();
  fut.get();
  auto end_time = std::chrono::steady_clock::now();

  std::chrono::duration<double, std::milli> execution_time =
      end_time - start_time;
  double milliseconds = execution_time.count();
  std::cout << milliseconds << "\n";

  EXPECT_EQ(ret->size(), 8'000);
  EXPECT_EQ(ret->at(0).size(), 3);
  scheduler.export_analyse("data3.bin");
  // NOLINTEND
}

TEST(SchedulerTest, HashJoinWithVector) {
  using namespace hashjoin;
  // build query plan
  Plan plan;

  // NOLINTBEGIN
  Table table1 = {
      {1, std::make_shared<std::vector<size_t>>(std::vector<size_t>{1, 2, 3}),
       4},
      {2, std::make_shared<std::vector<size_t>>(std::vector<size_t>{4, 5, 6}),
       5},
      {3, std::make_shared<std::vector<size_t>>(std::vector<size_t>{7, 8, 9}),
       6}};
  // NOLINTEND
  Table table2 = {{1, 1.1F},   // NOLINT
                  {2, 2.2F},   // NOLINT
                  {4, 3.3F}};  // NOLINT

  print_table(table1);
  print_table(table2);

  plan.new_input(std::move(table1));
  plan.new_input(std::move(table2));

  // projection of table1
  std::vector<std::tuple<size_t, DataType>> scan1_attrs = {
      {0, DataType::INT32}, {1, DataType::VECTOR}};  // get col 0 and col 1
  size_t scan1 = plan.new_scan_node(0, scan1_attrs);

  // projection of table2
  std::vector<std::tuple<size_t, DataType>> scan2_attrs = {
      {0, DataType::INT32}, {1, DataType::FLOAT}};  // get col 0 and col 1
  size_t scan2 = plan.new_scan_node(1, scan2_attrs);

  // hash join of table1 and table2
  std::vector<std::tuple<size_t, DataType>> join_attrs = {
      {0, DataType::INT32}, {1, DataType::VECTOR}, {3, DataType::FLOAT}};
  size_t join = plan.new_join_node(scan1, scan2, 0, 0, join_attrs);

  plan.root = join;

  // execute with scheduler
  Scheduler scheduler(8, 1);  // NOLINT
  scheduler.start();

  auto rg = std::make_unique<ResourceGroup>();
  auto ret = generate_resource_group(rg.get(), plan, plan.root);

  auto gpu_task = make_gpu_task_set(add_vectors_by_gpu, ret, 1);
  auto sum_vector = gpu_task->get_future();
  rg->add_task_set(std::move(gpu_task));

  auto finish = rg->get_future();
  // submit and wait for finishing
  scheduler.submit_resource_group(std::move(rg));
  finish.get();

  auto vector_result = sum_vector.get();

  print_table(*ret);
  std::vector<size_t> ground_truth = {5, 7, 9};  // NOLINT
  EXPECT_EQ(*vector_result, ground_truth);
}

// TEST(SchedulerTest, HashJoinWithLargeVector) {
//   using namespace hashjoin;
//   // build query plan
//   Plan plan;
//
//   // 生成100行数据，每个vector包含1亿条数据
//   const size_t ROW_COUNT = 100;
//   const size_t VECTOR_SIZE = 100'0000;  // 1亿
//
//   // 构建table1，包含100行，每行有一个大型vector
//   Table table1;
//   for (size_t i = 0; i < ROW_COUNT; i++) {
//     // 创建一个包含1亿条数据的vector
//     auto large_vector = std::make_shared<std::vector<size_t>>();
//     large_vector->reserve(VECTOR_SIZE);
//     for (size_t j = 0; j < VECTOR_SIZE; j++) {
//       large_vector->push_back(i * 10 + j % 10);  // 生成一些有规律的大数据
//     }
//
//     // 添加到表中
//     table1.push_back({static_cast<int32_t>(i + 1), std::move(large_vector),
//                       static_cast<int32_t>(i + 4)});
//   }
//
//   // 构建table2，包含100行
//   Table table2;
//   for (size_t i = 0; i < ROW_COUNT; i++) {
//     // 创建一些不匹配项和匹配项，让join有选择性
//     if (i % 3 == 0) {  // 只有id能被3整除的行能与table1连接
//       table2.push_back(
//           {static_cast<int32_t>(i + 1), static_cast<float>((i + 1) * 1.1)});
//     }
//   }
//
//   std::cout << "Table 1 size: " << table1.size() << " rows" << std::endl;
//   std::cout << "Table 2 size: " << table2.size() << " rows" << std::endl;
//   std::cout << "Vector size in each row: " << VECTOR_SIZE << " elements"
//             << std::endl;
//
//   plan.new_input(std::move(table1));
//   plan.new_input(std::move(table2));
//
//   // projection of table1
//   std::vector<std::tuple<size_t, DataType>> scan1_attrs = {
//       {0, DataType::INT32}, {1, DataType::VECTOR}};  // get col 0 and col 1
//   size_t scan1 = plan.new_scan_node(0, scan1_attrs);
//
//   // projection of table2
//   std::vector<std::tuple<size_t, DataType>> scan2_attrs = {
//       {0, DataType::INT32}, {1, DataType::FLOAT}};  // get col 0 and col 1
//   size_t scan2 = plan.new_scan_node(1, scan2_attrs);
//
//   // hash join of table1 and table2
//   std::vector<std::tuple<size_t, DataType>> join_attrs = {
//       {0, DataType::INT32}, {1, DataType::VECTOR}, {3, DataType::FLOAT}};
//   size_t join = plan.new_join_node(scan1, scan2, 0, 0, join_attrs);
//
//   plan.root = join;
//
//   // 使用更多的线程来处理大量数据
//   Scheduler scheduler(8, 1);  // 16线程，4个GPU任务队列
//   scheduler.start();
//
//   auto rg = std::make_unique<ResourceGroup>();
//   auto ret = generate_resource_group(rg.get(), plan, plan.root);
//
//   auto gpu_task = make_gpu_task_set(add_vectors_by_gpu, ret, 1);
//   auto sum_vector = gpu_task->get_future();
//   rg->add_task_set(std::move(gpu_task));
//
//   auto start_time = std::chrono::high_resolution_clock::now();
//
//   auto finish = rg->get_future();
//   // submit and wait for finishing
//   scheduler.submit_resource_group(std::move(rg));
//   finish.get();
//
//   auto end_time = std::chrono::high_resolution_clock::now();
//   auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(
//                       end_time - start_time)
//                       .count();
//   std::cout << "Execution time: " << duration << " ms" << std::endl;
//
//   auto vector_result = sum_vector.get();
//
//   // 如果vector_result非常大，只验证部分元素
//   if (vector_result && !vector_result->empty()) {
//     std::cout << "Vector result size: " << vector_result->size() << std::endl;
//     std::cout << "First 10 elements of result vector: ";
//     for (size_t i = 0; i < std::min(size_t(10), vector_result->size()); i++) {
//       std::cout << (*vector_result)[i] << " ";
//     }
//     std::cout << std::endl;
//
//     // // 验证部分结果
//     // // 实际验证会根据实际数据生成逻辑来确定预期的结果
//     // // 这里只是一个示例，实际需要根据数据生成逻辑调整
//     // std::vector<size_t> expected_first_elements = {
//     //     10, 20, 30};  // 替换为根据数据生成逻辑计算出的预期值
//     //
//     // for (size_t i = 0;
//     //      i < std::min(expected_first_elements.size(), vector_result->size());
//     //      i++) {
//     //   EXPECT_EQ((*vector_result)[i], expected_first_elements[i]);
//     // }
//   } else {
//     FAIL() << "Result vector is empty or null";
//   }
//
//   scheduler.export_analyse("data4.bin");
// }
